from pydantic import BaseModel
from typing import List, Optional, Literal
from datetime import datetime

class Message(BaseModel):
    role: Literal["user", "assistant"]
    content: str
    message_order: int
    token_count: Optional[int] = None
    model_name: Optional[str] = None
    processing_time_ms: Optional[int] = None
    created_at: Optional[datetime] = None

class Conversation(BaseModel):
    id: Optional[int] = None
    user_id: int
    title: str
    description: Optional[str] = None
    status: Literal["active", "archived", "deleted"] = "active"
    message_count: int = 0
    created_at: Optional[datetime] = None
    updated_at: Optional[datetime] = None

class ConversationCreate(BaseModel):
    user_id: int
    title: str
    description: Optional[str] = None

class ConversationUpdate(BaseModel):
    title: Optional[str] = None
    description: Optional[str] = None
    status: Optional[Literal["active", "archived", "deleted"]] = None

class MessageCreate(BaseModel):
    conversation_id: int
    role: Literal["user", "assistant"]
    content: str
    message_order: int

class MessageUpdate(BaseModel):
    content: Optional[str] = None
    token_count: Optional[int] = None
    model_name: Optional[str] = None
    processing_time_ms: Optional[int] = None

class ConversationResponse(BaseModel):
    conversation: Conversation
    messages: List[Message]

class StreamResponse(BaseModel):
    type: Literal["message", "error", "complete"]
    content: str
    message_order: Optional[int] = None
    token_count: Optional[int] = None
    model_name: Optional[str] = None
    processing_time_ms: Optional[int] = None

class ChatRequest(BaseModel):
    """Unified request model for chat endpoint"""
    user_id: int
    message: str
    conversation_id: Optional[int] = None
    title: Optional[str] = None
    description: Optional[str] = None
    message_order: Optional[int] = None
    image_base64: Optional[str] = None  # 可选的图片base64编码
